import json
from datetime import datetime

import numpy as np
from typing import Dict, List


class TrafficGenerator:


# ...（保持原有类定义不变，参考用户提供的代码）...
    def __init__(self,
                 num_nodes: int = 4, # 定义节点数量
                 # min_value: float = 500,
                 # max_value: float = 800,
                 burst_config: Dict = None):
        self.num_nodes = num_nodes
        self.tor_nodes = [f"ToR{i+1}" for i in range(num_nodes)]  # 生成ToR名称列表（从1到4）
        self.time = 0.0
        # self.burst_intensity = 800

        # 突发流量配置
        self.burst_config = burst_config or {
            'hotspots': [0, 3],  # 突发热点节点
            # 'burst_duration': 0.5,  # 突发持续时间(s)
            # 'lambda_poisson': 0.3,  # 泊松分布参数
            'burst_intensity': 100.0  # 突发流量强度
        }

    def generate_all_to_all(self,
                            min_value: float = 20,
                            max_value: float = 100) -> np.ndarray:
        """
        生成可配置范围的All-to-All流量矩阵

        参数:
            min_value (float): 最小流量值，默认500GB
            max_value (float): 最大流量值，默认800GB

        返回:
            np.ndarray: N x N的流量矩阵，对角线置零
        """
        # 参数有效性校验
        if min_value < 0 or max_value < 0:
            raise ValueError("流量值不能为负数")
        if min_value >= max_value:
            raise ValueError("最大流量值必须大于最小流量值")

        # 生成均匀分布流量矩阵
        matrix = np.random.uniform(min_value, max_value,
                                   (self.num_nodes, self.num_nodes))
        np.fill_diagonal(matrix, 0)  # 对角线置零
        return matrix

    def generate_ring_allreduce(self,
                            min_value: float = 100,
                            max_value: float = 300) -> np.ndarray:
        """生成环形AllReduce流量矩阵"""
        matrix = np.zeros((self.num_nodes, self.num_nodes))
        chunk_size = np.random.uniform(min_value, max_value) # 定义每个节点的流量大小

        # 环形通信模式
        for i in range(self.num_nodes):
            # 顺时针方向通信
            src = i
            dst = (i + 1) % self.num_nodes
            matrix[src, dst] = chunk_size

            # 逆时针方向通信
            src = (i - 1) % self.num_nodes
            dst = i
            matrix[src, dst] = chunk_size

        return matrix

    def generate_burst_traffic(self, burst_intensity: float = 800) -> np.ndarray:
        """生成突发流量矩阵"""
        matrix = np.zeros((self.num_nodes, self.num_nodes))

        # 突发热点通信
        for src in self.burst_config['hotspots']:
            for dst in range(self.num_nodes):
                if src != dst:
                    matrix[src][dst] += np.random.exponential(
                        # self.burst_config['burst_intensity']
                        burst_intensity
                    )

        # 泊松过程控制突发次数
        # if np.random.poisson(self.burst_config['lambda_poisson']) > 0:
        hotspot = np.random.choice(self.burst_config['hotspots'])
        for _ in range(int(np.random.exponential(5))):
            dst = np.random.choice(self.num_nodes)
            if hotspot != dst:
                matrix[hotspot][dst] += np.random.exponential(
                    self.burst_config['burst_intensity'] * 5
                )

        return matrix

def generate_traffic_records(
        num_nodes: int = 4,
        base_scale_factors: List[float] = [1.0, 1.5 ,2.0,2.5, 3.0,3.5 ,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0],
        # burst_intensity_factors: List[float] = [2.0,3.0,5.0]
) -> Dict:
    """
    生成包含多种流量模式的流量记录

    参数：
        num_nodes: 网络节点数量
        base_scale_factors: 基础流量缩放系数列表
        burst_intensity_factors: 突发强度系数列表

    返回：
        包含所有流量记录的字典
    """
    records = {
        "all_to_all": [],
        "all_to_all_burst": [],
        "ring_allreduce": [],
        "ring_allreduce_burst": []
    }

    generator = TrafficGenerator(num_nodes=num_nodes)

    # 生成基础流量模式
    for scale in base_scale_factors:
        # All-to-All
        base_min, base_max = 40 * scale, 60 * scale
        print(f"获得的all-to-all 的缩放比例为 {scale}")
        all_to_all = generator.generate_all_to_all(base_min, base_max)
        # burst_all_min,burst_all_max=base_min *5 ,base_max *5
        # burst_all=generator.generate_all_to_all(burst_all_min,burst_all_max)
        burst_all=generator.generate_burst_traffic(base_max*2)
        all_to_all_burst=all_to_all+burst_all

        records["all_to_all"].append({
            "scale_factor": scale,
            "min_value": base_min,
            "max_value": base_max,
            "matrix": all_to_all.tolist()
        })

        # Ring AllReduce
        base_min_1, base_max_1 = 200 * scale, 300* scale
        ring = generator.generate_ring_allreduce(base_min_1, base_max_1)
        # burst_min_ring, burst_max_ring=base_min_1*5,base_max_1 *5
        # burst_ring = generator.generate_all_to_all(burst_min_ring, burst_max_ring)
        burst_ring=generator.generate_burst_traffic(base_max_1*5)
        # burst_all_, burst_all_max = base_min_1 * 5, base_max_1 * 5
        # burst_ring=generator.generate_all_to_all(burst_all_min, burst_all_max)
        # ring_burst=ring+burst_all

        records["ring_allreduce"].append({
            "scale_factor": scale,
            "min_value": base_min,
            "max_value": base_max,
            "matrix": ring.tolist()
        })

        # 生成带突发的流量模式
        # for burst_scale in burst_intensity_factors:
        #     # All-to-All + Burst
        #     # burst_intensity = 100.0 * scale * burst_scale
        #     burst_intensity = round(base_max * 2 * burst_scale)
        #     burst_matrix = generator.generate_burst_traffic(burst_intensity)
        #     combined = all_to_all + burst_matrix
        #     np.fill_diagonal(combined, 0)  # 保持对角线为0

        # 生成带突发的流量形式

        records["all_to_all_burst"].append({
            "base_scale": scale,
            "burst_factor": scale*5,
            # "min_value": burst_all_min,
            # "max_value": burst_all_max,
            # "total_intensity": burst_intensity,
            "matrix": all_to_all_burst.tolist()
        })

        # Ring AllReduce + Burst
        ring_burst = ring + burst_ring
        np.fill_diagonal(ring_burst, 0)

        records["ring_allreduce_burst"].append({
            "base_scale": scale,
            "burst_factor": scale * 5,
            # "min_value": burst_min_ring,
            # "max_value": burst_max_ring,
            # "scale_factor": burst_scale,
            # "total_intensity": burst_intensity,
            "matrix": ring_burst.tolist()
        })

    return records


# 生成并保存流量记录
if __name__ == "__main__":
    base_scale_factors = []
    burst_intensity_factors=[]
    # base_scale_factors=[ 0.5, 1.0, 2.0, 3.0, 5.0, 8.0, 10.0,15.0,20.0,25,0,30.0]
    # 迭代50次
    num_iter1=50
    # num_iter2 = 25
    scale_base=0.0
    # scale_burst=0.0
    for i in range(num_iter1):
    # for i in len(base_scale_factors):
        print("scale_base第{}次迭代".format(i+1))
        scale_base=scale_base+5
        # scale_base=base_scale_factors[i]
        base_scale_factors.append(scale_base)
    # for j in range(num_iter2):
    #     print("scale_burst第{}次迭代".format(j+1))
    #     scale_burst = scale_burst + 0.5
    #     burst_intensity_factors.append(scale_burst)
    # print(base_scale_factors)
    # print(burst_intensity_factors)
    # # 参数设置
    config = {
        "num_nodes": 16,
        # "base_scale_factors": [1.0, 1.5 ,2.0,2.5, 3.0,3.5 ,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0],  # 基础流量从50%到200%
        # "burst_intensity_factors": [1.0, 1.5 ,2.0,2.5, 3.0,3.5 ,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0]  # 突发强度从5x到15x
        "base_scale_factors": base_scale_factors,
        # "burst_intensity_factors": burst_intensity_factors
    }

    # 生成数据
    traffic_data = generate_traffic_records(**config)

    # 添加元数据
    metadata = {
        "description": "网络流量模拟数据集 (单位: GB)",
        "node_naming": "ToR1-ToR4 表示4个Top of Rack交换机",
        "matrix_interpretation": "matrix[i][j] 表示从ToR(i+1)到ToR(j+1)的流量",
        "timestamp": datetime.now().isoformat()
    }
    traffic_data["metadata"] = metadata

    # 保存为JSON文件
    with open("traffic_evolution_16.json", "w") as f:
        json.dump(traffic_data, f, indent=2)

    print("流量数据已保存至 traffic_evolution_16.json")

import matplotlib.pyplot as plt



